'''
Created on Apr 18, 2013

@author: LarsoMi1
'''

import sys


import matplotlib.pyplot as plt
from matplotlib.lines import Line2D
from matplotlib.ticker import FuncFormatter

from data_types.InternetUsageSurvey import income_levels

_population_by_age = {}

def initialize_population(session):
    if not _population_by_age:
        for household in session.households.values():
            for person in household:
                _population_by_age[person.age] = _population_by_age.get(person.age, 0) + 1

def graph_session(session, verbose=False):
    people = [p for h in session.households.values() for p in h]
    initialize_population(session)

    _graph_internet_usage_by_age(session.households.values(), verbose)
    _graph_internet_usage_for_education_by_age(session.households.values(), people)
    _graph_luddites_by_age(session.households.values(), verbose)
    _graph_internet_usage_by_education_level(session.households.values(), verbose)
    _graph_internet_usage_by_state(session.households.values(), verbose)
    _graph_incomes(session.households.values(), verbose)
    try:
        plt.show()
    except AttributeError, error:
        sys.stderr.write('could not display plot')
        if 0 < verbose:
            sys.stderr.write(' (%s)' % str(error))
        sys.stderr.write('\n')

def _graph_luddites_by_age(households, verbose=False):
    fig = plt.figure()
    fig.canvas.set_window_title('Luddites')
    
    luddites = {}
    users = {}
    
    for household in households:
        if household.home_internet_access != 'Y':
            continue
        for person in household:
            if person.home_internet_access == 'N':
                luddites[person.age] = luddites.get(person.age, 0) + 1
            else:
                users[person.age] = users.get(person.age, 0) + 1
    
    ages = range(5, 80)
    totals = {}
    for age in ages:
        totals[age] = users[age] + luddites[age]
    user_percentages = list(users[i] / float(totals[i]) for i in totals.iterkeys())
    luddite_percentages = list(luddites[i] / float(totals[i]) for i in totals.iterkeys())
    
    sp = fig.add_subplot(211)
    rects1 = sp.bar(ages, [luddites[i] for i in ages], color='r')
    plt.xlabel('Age')
    plt.xlim((ages[0], ages[-1]))
    plt.ylabel('Population (People)')
    plt.title('Luddites (totals)')
    plt.legend([rects1[0]], ['Luddites'], fontsize='small')
    
    sp = fig.add_subplot(212)
    
    rects2 = sp.bar(ages, luddite_percentages, width=1.0, color='r')
    rects3 = sp.bar(ages, user_percentages, bottom=luddite_percentages, width=1.0, color='b')
    plt.xlabel('Age')
    plt.xlim((ages[0], ages[-1]))
    plt.ylabel('Percentage of population')
    plt.title('Luddites (percentages)')
    plt.legend((rects2[0], rects3[0]), ('Luddites', 'Users'), fontsize='small')
    plt.gca().yaxis.set_major_formatter(FuncFormatter(_ticks_percentage))
    
    fig.tight_layout()


def _get_groups_of_people(households, discriminator='age'):
    haves = {}
    have_nots = {}
    unknowns = []
    for household in households:
        for person in household:
            if person.home_internet_access == 'Y':
                if not haves.has_key(getattr(person, discriminator)):
                    haves[getattr(person, discriminator)] = []
                haves[getattr(person, discriminator)].append(person)
            elif person.home_internet_access == 'N':
                if not have_nots.has_key(getattr(person, discriminator)):
                    have_nots[getattr(person, discriminator)] = []
                have_nots[getattr(person, discriminator)].append(person)
            else:
                unknowns.append(person)
    return (haves, have_nots, unknowns)

def _graph_internet_usage_by_age(households, verbose=False):
    fig = plt.figure()
    fig.canvas.set_window_title('Internet Usage by Age')
    
    haves, have_nots = _get_groups_of_people(households)[:2]
    
    ages = range(5, 79)
    populations = {
        'haves': dict(zip(ages, [0] * len(ages))),
        'phaves': dict(zip(ages, [0] * len(ages))),
        'have_nots': dict(zip(ages, [0] * len(ages))),
        'phave_nots': dict(zip(ages, [0] * len(ages)))
    }
    
    for (age, people) in sorted(haves.iteritems()):
        if ages[0] <= age <= ages[-1]:
            populations['haves'][age] = len(people)
        elif verbose:
            print age, '--', len(people)
    
    for (age, people) in sorted(have_nots.iteritems()):
        if ages[0] <= age <= ages[-1]:
            populations['have_nots'][age] = len(people)
        elif verbose:
            print age, '--', len(people)
    
    for age in ages:
        total_population = float(populations['haves'].get(age, 0) + populations['have_nots'].get(age, 0))
        populations['phaves'][age] = populations['haves'].get(age, 0) / total_population
        populations['phave_nots'][age] = populations['have_nots'].get(age, 0) / total_population
    phaves = list(item[-1] for item in sorted(populations['phaves'].iteritems()))
    phave_nots = list(item[-1] for item in sorted(populations['phave_nots'].iteritems()))
    
    sp = fig.add_subplot(111)
    rects1 = sp.bar(ages, phaves, width=1.0, color='r')
    rects2 = sp.bar(ages, phave_nots, bottom=phaves, width=1.0, color='b')
    
    sp.add_line(Line2D([0, 79], [0.5, 0.5], transform=sp.transAxes, linewidth=1, color='k'))
    
    plt.xlabel('Age')
    plt.xlim([5, 79])
    plt.ylabel('Percentage of population')
    plt.title('Internet Access by Age')
    plt.legend((rects1[0], rects2[0]), ('Home internet access', 'No home internet access'))
    plt.gca().yaxis.set_major_formatter(FuncFormatter(_ticks_percentage))
    
    fig.tight_layout()

def _graph_incomes(households, verbose=False):
    fig = plt.figure()
    fig.canvas.set_window_title('Household Income by State')
    
    incomes_by_state = {}
    for household in households:
        if not incomes_by_state.has_key(household.state):
            incomes_by_state[household.state] = []
        incomes_by_state[household.state].append(household.HEFAMINC)
    
    data = [incomes_by_state['CO'], incomes_by_state['KS'], incomes_by_state['OK']]
    fig.add_subplot(111)
    plt.boxplot(data)
    
    plt.xlabel('State')
    plt.xticks([1, 2, 3], ['CO', 'KS', 'OK'])
    plt.ylim([0, 17])
    plt.ylabel('Income')
    plt.yticks(range(1, 17), [income_levels[i] for i in range(1, 17)])
    plt.title('Household Income by State')
    
    fig.tight_layout()

def _graph_internet_usage_by_state(households, verbose=False):
    fig = plt.figure()
    fig.canvas.set_window_title('Internet Usage by State')
    
    haves, have_nots = _get_groups_of_people(households, 'state')[:2]
    haves_populations =     list(len(    haves[state]) for state in haves.iterkeys())
    have_nots_populations = list(len(have_nots[state]) for state in haves.iterkeys())
    
    if verbose:
        ks   = {'haves': haves['KS'], 'have_nots': have_nots['KS'], 'total': float(len(haves['KS'] + have_nots['KS']))}
        ok = {'haves': haves['OK'], 'have_nots': have_nots['OK'], 'total': float(len(haves['OK'] + have_nots['OK']))}
        co = {'haves': haves['CO'], 'have_nots': have_nots['CO'], 'total': float(len(haves['CO'] + have_nots['CO']))}
        print 'KS:    access : %4d (%1.2f)' % (len(ks['haves']),     len(ks['haves'])     / ks['total'])
        print 'KS: no access : %4d (%1.2f)' % (len(ks['have_nots']), len(ks['have_nots']) / ks['total'])
        print 'OK:    access : %4d (%1.2f)' % (len(ok['haves']),     len(ok['haves'])     / ok['total'])
        print 'OK: no access : %4d (%1.2f)' % (len(ok['have_nots']), len(ok['have_nots']) / ok['total'])
        print 'CO:    access : %4d (%1.2f)' % (len(co['haves']),     len(co['haves'])     / co['total'])
        print 'CO: no access : %4d (%1.2f)' % (len(co['have_nots']), len(co['have_nots']) / co['total'])
    
    relevant_populations = [float(haves_populations[i] + have_nots_populations[i]) for i in range(len(haves_populations))]
    
    haves_percentages = [haves_populations[i] / relevant_populations[i] for i in range(len(relevant_populations))]
    have_nots_percentages = [have_nots_populations[i] / relevant_populations[i] for i in range(len(relevant_populations))]
    
    keys = range(len(haves_populations))
    
    sp = fig.add_subplot(211)
    rects1 = sp.bar(keys,                        haves_populations, width=0.5, color='r')
    rects2 = sp.bar([x + 0.5 for x in keys], have_nots_populations, width=0.5, color='b')
    plt.xlabel('State')
    plt.xticks([x + 0.5 for x in keys], list(haves.iterkeys()), rotation='vertical', size='small')
    plt.ylabel('Population Size (people)')
    plt.title('Internet Access by State')
    plt.legend((rects1[0], rects2[0]), ('Home internet access', 'No home internet access'), fontsize='small')
    
    sp = fig.add_subplot(212)
    rects3 = sp.bar(keys,     haves_percentages,                           width=1.0, color='r')
    rects4 = sp.bar(keys, have_nots_percentages, bottom=haves_percentages, width=1.0, color='b')
    
    sp.add_line(Line2D([0, len(keys)], [0.5, 0.5], transform=sp.transAxes, linewidth=2, color='k'))
    
    plt.xlabel('State')
    plt.xticks([x + 0.5 for x in keys], list(haves.iterkeys()), rotation='vertical', size='small')
    plt.ylabel('Percentage of population')
    plt.title('Internet Access by State')
    plt.legend((rects3[0], rects4[0]), ('Home internet access', 'No home internet access'), loc=3, fontsize='small')
    plt.gca().yaxis.set_major_formatter(FuncFormatter(_ticks_percentage))
    
    fig.tight_layout()

def _graph_internet_usage_by_education_level(households, verbose=False):
    fig = plt.figure()
    fig.canvas.set_window_title('Internet Usage by Education')
    
    haves, have_nots = _get_groups_of_people(households, 'education')[:2]
    
    levels = {
         0: 'LESS THAN 1ST GRADE',
         1: '1ST, 2ND, 3RD OR 4TH GRADE',
         2: '5TH OR 6TH GRADE',
         3: '7TH OR 8TH GRADE',
         4: '9TH GRADE',
         5: '10TH GRADE',
         6: '11TH GRADE',
         7: '12TH GRADE NO DIPLOMA',
         8: 'HIGH SCHOOL GRAD-DIPLOMA OR EQUIV (GED)',
         9: 'SOME COLLEGE BUT NO DEGREE',
        10: 'ASSOCIATE DEGREE-OCCUPATIONAL/VOCATIONAL',
        11: 'ASSOCIATE DEGREE-ACADEMIC PROGRAM',
        12: "BACHELOR'S DEGREE",       #(EX: BA, AB, BS)
        13: "MASTER'S DEGREE",         #(EX: MA, MS, MEng, MEd, MSW)
        14: 'PROFESSIONAL SCHOOL DEG', #(EX: MD, DDS, DVM)'
        15: 'DOCTORATE DEGREE'         #(EX: PhD, EdD)
    }
    
    keys = range(len(levels.keys()))
    
    haves_populations = list(len(        haves[item[-1]]) for item in sorted(levels.iteritems()))
    have_nots_populations = list(len(have_nots[item[-1]]) for item in sorted(levels.iteritems()))
    
    relevant_populations = [float(haves_populations[i] + have_nots_populations[i]) for i in range(len(haves_populations))]
    
    haves_percentages = [haves_populations[i] / relevant_populations[i] for i in range(len(relevant_populations))]
    have_nots_percentages = [have_nots_populations[i] / relevant_populations[i] for i in range(len(relevant_populations))]
    
    sp = fig.add_subplot(111)
    rects1 = sp.bar(keys,     haves_percentages,                           width=1.0, color='r')
    rects2 = sp.bar(keys, have_nots_percentages, bottom=haves_percentages, width=1.0, color='b')
    
    plt.xlabel('Education Level')
    plt.xticks([x + 0.5 for x in keys], [item[-1].lower() for item in sorted(levels.iteritems())], rotation='vertical', size='x-small')
    plt.ylabel('Percentage of Population')
    plt.title('Internet Access by Education Level')
    plt.legend((rects1[0], rects2[0]), ('Home internet access', 'No home internet access'), loc=2)
    plt.gca().yaxis.set_major_formatter(FuncFormatter(_ticks_percentage))
    
    fig.tight_layout()

def _graph_internet_usage_for_education_by_age(households, people, verbose=False):
    fig = plt.figure()
    fig.canvas.set_window_title('Internet Usage for Education by Age')
    
    ages = range(5, 80)
    relevant_people = {}
    users = {}
    non_users = {}
    unknowns = []
    
    for person in people:
        if person.internet_access == 'Y' and ages[0] <= person.age <= ages[-1]:
            if not relevant_people.has_key(person.age):
                relevant_people[person.age] = []
            if person.educational_internet_usage == 'Y':
                if not users.has_key(person.age):
                    users[person.age] = []
                users[person.age].append(person)
                relevant_people[person.age].append(person)
            elif person.educational_internet_usage == 'N':
                if not non_users.has_key(person.age):
                    non_users[person.age] = []
                non_users[person.age].append(person)
                relevant_people[person.age].append(person)
            else:
                unknowns.append(person)
    
    num_users = sum(len(age_group) for age_group in users.itervalues())
    num_non_users = sum(len(age_group) for age_group in non_users.itervalues())
    num_relevant = sum(len(age_group) for age_group in relevant_people.itervalues())
    
    if verbose:
        print 'users:     %5d' % num_users
        print 'non-users: %5d' % num_non_users
        print 'relevant:  %5d (%5d)' % (num_relevant, num_users + num_non_users)
        print 'unknown:   %5d' % len(unknowns)
        print 'total:     %5d' % (num_users + num_non_users + len(unknowns))
    
    populations = {
        'users': dict(zip(ages, [0] * len(ages))),
        'non_users': dict(zip(ages, [0] * len(ages))),
        'relevant': dict(zip(ages, [0] * len(ages)))
    }
    
    for (age, people) in sorted(relevant_people.iteritems()):
        if ages[0] <= age <= ages[-1]:
            populations['relevant'][age] = len(people)
        elif verbose:
            print age, '--', len(people)
    relevant_populations = list(pair[-1] for pair in sorted(populations['relevant'].iteritems()))
    
    for (age, people) in sorted(users.iteritems()):
        if ages[0] <= age <= ages[-1]:
            populations['users'][age] = len(people)
        elif verbose:
            print age, '--', len(people)
    user_populations = list(pair[-1] for pair in sorted(populations['users'].iteritems()))
    
    for (age, people) in sorted(non_users.iteritems()):
        if ages[0] <= age <= ages[-1]:
            populations['non_users'][age] = len(people)
        elif verbose:
            print age, '--', len(people)
    non_user_populations = list(pair[-1] for pair in sorted(populations['non_users'].iteritems()))
    
    for index in range(len(ages)):
        if relevant_populations[index]:
            user_populations[index] /= float(relevant_populations[index])
            non_user_populations[index] /= float(relevant_populations[index])
    
    sp = fig.add_subplot(111)
    
    rects3 = sp.bar(ages, user_populations,                              width=1.0, color='r')
    rects4 = sp.bar(ages, non_user_populations, bottom=user_populations, width=1.0, color='b')
    
    plt.xlabel('Age')
    plt.xlim([15, ages[-1]])
    plt.ylabel('Percentage of population')
    plt.title('Educational Usage: Stacked Bar Graph')
    plt.legend((rects3[0], rects4[0]), ('Educational Usage', 'No Educational Usage'), fontsize='small')
    plt.gca().yaxis.set_major_formatter(FuncFormatter(_ticks_percentage))
    
    fig.tight_layout()

def _ticks_percentage(value, position):
    # Ignore the passed in position. This has the effect of scaling the default
    # tick locations.
    s = str(100 * value)

    # The percent symbol needs escaping in latex
    if plt.rcParams['text.usetex'] == True:
        return s + r'$\%$'
    else:
        return s + '%'
